import os
import subprocess
import logging
from pathlib import Path
import wave
import json
import librosa
import numpy as np
import soundfile as sf

logger = logging.getLogger(__name__)


class NoAudioError(Exception):
    """输入文件没有音频流"""
    pass


def has_audio_stream(input_path: str) -> bool:
    """
    使用 ffprobe 检查输入文件是否包含音频流。
    需要系统安装 ffprobe，并且在 PATH 中可调用。
    """
    cmd = [
        "ffprobe",
        "-v", "error",
        "-select_streams", "a",
        "-show_entries", "stream=index",
        "-of", "json",
        input_path
    ]
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        # 可能文件不存在或格式错，我们也当作“无音频”
        return False
    info = json.loads(p.stdout or "{}")
    streams = info.get("streams", [])
    return len(streams) > 0


def ensure_wav(input_path: str, output_wav: str, target_sr: int = 16000, mono: bool = True):
    """
    使用 FFmpeg 将任意音频/视频 转换成标准的 WAV：
      - PCM s16le 编码
      - 单声道（mono）
      - 采样率 target_sr（默认 16000）
    input_path: 原始文件（可以是 .mp4/.m4a/.aac/.wav/其他）
    output_wav: 转换后输出路径 (扩展名应为 .wav)
    """
    # 确保父目录存在
    out_dir = Path(output_wav).parent
    out_dir.mkdir(parents=True, exist_ok=True)

    if not has_audio_stream(input_path):
        raise NoAudioError(f"{input_path} 中未检测到音频流")

    output_path = Path(output_wav)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # 构造 FFmpeg 命令
    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-acodec", "pcm_s16le",
        "-ac", "1" if mono else "2",
        "-ar", str(target_sr),
        output_wav
    ]
    try:
        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except subprocess.CalledProcessError as e:
        logger.error(f"FFmpeg 转码失败: {e}")
        raise RuntimeError(f"ensure_wav: FFmpeg 转码失败, cmd: {' '.join(cmd)}")
    return output_wav


def wav2pcm(wav_file: str, pcm_file: str, data_type=np.int16):
    """
    将标准 WAV (16k PCM mono) 转成原始 PCM 文件 (s16le)。
    只适用于输入本身就是 PCM s16le 的 WAV。
    data_type: numpy dtype，默认为 np.int16，对应 s16le。
    """
    # 确保输出目录存在
    os.makedirs(os.path.dirname(pcm_file), exist_ok=True)

    with open(wav_file, "rb") as f_in:
        # 跳过 WAV 文件头（44 字节）
        f_in.seek(44)
        pcm_data = np.fromfile(f_in, dtype=data_type)

    # 将 PCM 数据写入
    pcm_data.tofile(pcm_file)
    return pcm_file


def pcm2wav(pcm_file: str, wav_file: str, channels: int = 1, bits: int = 16, sample_rate: int = 16000):
    """
    将原始 PCM（s16le）转成带 WAV 头的标准 WAV：
      channels: 通道数 (1 表示单声道)
      bits: 位深 (通常是 16)
      sample_rate: 采样率 (一般 16000)
    """
    if bits % 8 != 0:
        raise ValueError(f"bits 必须是 8 的倍数，目前 bits={bits}")

    os.makedirs(os.path.dirname(wav_file), exist_ok=True)
    with wave.open(wav_file, "wb") as wav_out:
        wav_out.setnchannels(channels)
        wav_out.setsampwidth(bits // 8)
        wav_out.setframerate(sample_rate)
        # 读取所有 PCM 数据
        with open(pcm_file, "rb") as f_pcm:
            pcm_bytes = f_pcm.read()
            wav_out.writeframes(pcm_bytes)
    return wav_file


def split_wav(input_wav: str, segment_duration: float = 60.0,
              output_dir: str = None, to_pcm: bool = True):
    """
    将一个标准 WAV 文件按 segment_duration (秒) 切成若干段：
      - 先用 librosa.load 一次性加载音频数据
      - 按 segment_duration 划分 numpy 数组片段
      - 把每个片段写成 WAV，如果 to_pcm=True 会额外转成 PCM，并删除中间 WAV
    返回：每个片段的 PCM 文件路径列表（如果 to_pcm=False，则返回 WAV 路径列表）
    """
    # 使用 librosa 读取整个 WAV（它已是 16k 单声道 PCM）
    y, sr = librosa.load(input_wav, sr=None, mono=True)
    total_duration = len(y) / sr

    # 如果整个音频时长不超过 segment_duration，直接返回单个文件
    if total_duration <= segment_duration:
        if to_pcm:
            pcm_name = str(Path(output_dir or Path(input_wav).parent) /
                           f"{Path(input_wav).stem}.pcm")
            return [wav2pcm(input_wav, pcm_name)]
        else:
            return [input_wav]

    # 计算要切成多少段
    num_segments = int(np.ceil(total_duration / segment_duration))
    parent_dir = Path(output_dir or Path(input_wav).parent / "segments")
    parent_dir.mkdir(parents=True, exist_ok=True)

    pcm_paths = []
    for i in range(num_segments):
        start_sample = int(i * segment_duration * sr)
        end_sample = int(min((i + 1) * segment_duration * sr, len(y)))
        segment_data = y[start_sample:end_sample]
        if segment_data.size == 0:
            continue

        # 先写中间 WAV
        seg_wav_path = str(parent_dir / f"{Path(input_wav).stem}_seg_{i}.wav")
        sf.write(seg_wav_path, segment_data, sr)

        if to_pcm:
            # 再把该段 WAV 转为 PCM 并删除中间 WAV
            seg_pcm_path = str(parent_dir / f"{Path(input_wav).stem}_seg_{i}.pcm")
            wav2pcm(seg_wav_path, seg_pcm_path)
            pcm_paths.append(seg_pcm_path)
            try:
                os.remove(seg_wav_path)
            except Exception as e:
                logger.warning(f"无法删除临时 WAV: {seg_wav_path}，{e}")
        else:
            pcm_paths.append(seg_wav_path)

    return pcm_paths


def split_wav_in_memory(input_wav: str, segment_duration: float = 60.0):
    """
    如果只想要各段的 numpy 数组而不写磁盘，可以调用此函数：
    返回值：[(segment_array, sr), ...] 列表
    """
    y, sr = librosa.load(input_wav, sr=None, mono=True)
    total_duration = len(y) / sr
    if total_duration <= segment_duration:
        return [(y, sr)]

    num_segments = int(np.ceil(total_duration / segment_duration))
    segments = []
    for i in range(num_segments):
        start_sample = int(i * segment_duration * sr)
        end_sample = int(min((i + 1) * segment_duration * sr, len(y)))
        seg_data = y[start_sample:end_sample]
        if seg_data.size == 0:
            continue
        segments.append((seg_data, sr))
    return segments


def process_full_audio(input_path: str, working_dir: str, segment_seconds: float = 60.0):
    """
    示例：从任意 input_path (视频或音频) → 转 WAV → 切分 → 得到各段 PCM
    返回：一个列表，包含所有 PCM 片段的绝对路径。
    1. ensure_wav: 转码成 16k 单声道 PCM WAV
    2. split_wav: 按 segment_seconds 切片，并输出 PCM
    """
    # 1. 转成标准 WAV
    wav_basename = Path(input_path).stem
    wav_output = str(Path(working_dir) / f"{wav_basename}_16k_mono.wav")
    ensure_wav(input_path, wav_output, target_sr=16000, mono=True)

    # 2. 切分并输出 PCM
    pcm_segments = split_wav(
        input_wav=wav_output,
        segment_duration=segment_seconds,
        output_dir=working_dir,
        to_pcm=True
    )

    return pcm_segments
